from uuid import uuid4
from pathlib import Path
from sanic import Request, Websocket
from sanic import Blueprint
from sanic.response import HTTPResponse
from sanic.log import logger
from aiohttp import ClientSession, ClientTimeout, WSMsgType

from apps.base.hunter import AppError
from apps.base.response import CustomerResponse
from modules.xfyun.spark import SparkAuthExt, send_data, rcev_data, process_file_and_send_to_spark
from modules.xfyun.chat import ChatAuthExt, send_chat, recv_chat
from modules.xfyun.tts import TtsAuthExt, send_tts, recv_tts
from modules.xfyun.draw import du_draw
from settings import STATIC_PATH
import aiohttp
import json
import base64
import aiofiles
message_bp = Blueprint("message", url_prefix="message", version=1)


# /v1/message/upload
# @message_bp.post("/stt")
# async def upload(request: Request) -> HTTPResponse:
#     """
#     文件上传
#     """
#     if "file" not in request.files:
#         raise AppError(10001, "the field: file is not found")
#
#     files = request.files['file']
#     # print(file)
#
#     # 强制修改文件名称, 防止文件重名
#     file_name = uuid4().hex
#
#     # 保存文件到本地
#     with open(Path.joinpath(STATIC_PATH, file_name), 'wb') as f:
#         f.write(files[0].body)
#
#     file_path = Path.joinpath(STATIC_PATH, file_name)
#     spark_data = await process_file_and_send_to_spark(file_path, logger)
#     if spark_data is not None:
#         # 处理返回的数据
#         print(f"Received data from Spark: {spark_data}")
#     else:
#         print("Failed to process file and send to Spark")
#     data = {
#         "file": spark_data,
#     }
#     return CustomerResponse(data=data)
#

@message_bp.websocket("/stt")
async def chat(request: Request, ws: Websocket) -> None:
    await ws.accept()

    try:
        while True:
            data = await ws.receive_bytes()
            # Handle received data, e.g., save to file
            file_name = uuid4().hex
            file_path = Path.joinpath(STATIC_PATH, file_name)

            with open(file_path, 'wb') as f:
                f.write(data)

            # Process file and send to Spark
            spark_data = await process_file_and_send_to_spark(file_path)

            # Send response back to client
            if spark_data is not None:
                await ws.send_text(f"Received data from Spark: {spark_data}")
            else:
                await ws.send_text("Failed to process file and send to Spark")

    except ws.exceptions.ConnectionClosedOK:
        print("WebSocket connection closed")


async def process_file_and_send_to_spark(file_path):
    # Implement your logic to process file and send to Spark
    pass

# /v1/message/chat
@message_bp.websocket("/chat")
async def chat(request: Request, ws: Websocket) -> None:
    """
    对接spark模型
    """
    # 接收对话参数
    query = await ws.recv()
    user_id = uuid4()

    async with ClientSession(timeout=ClientTimeout(total=30)) as session:
        # 初始化
        wsParam = ChatAuthExt()
        # 创建认证url
        url = wsParam.create_url()
        # 发起websocket连接
        async with session.ws_connect(url) as chat_ws:
            logger.info(F"user query: {query}")
            # 发送数据到spark
            await send_chat(chat_ws, wsParam, user_id, query)
            # 接收数据
            async for msg in chat_ws:
                if msg.type == WSMsgType.TEXT:
                    chat_data = await recv_chat(chat_ws, msg.data)
                    logger.info(F"spark_data: {chat_data}")
                    # 发送数据给ESP客户端
                    await ws.send(chat_data)


@message_bp.websocket("/tts")
async def tts(request: Request, ws: Websocket) -> HTTPResponse:
    try:
        query = await ws.recv()
        user_id = uuid4()

        async with ClientSession(timeout=ClientTimeout(total=30)) as session:
            # 初始化
            wsParam = TtsAuthExt(text=query)
            # 创建认证url
            url = wsParam.create_url()
            # 发起websocket连接
            async with session.ws_connect(url) as tts_ws:
                await send_tts(tts_ws, wsParam)

                # 接收数据
                async for msg in tts_ws:
                    if msg.type == aiohttp.WSMsgType.TEXT:
                        received_message = msg.data
                        print(f"Received message: {received_message}")

                        try:
                            audio = await recv_tts(tts_ws, received_message)
                            # 在这里处理接收到的音频数据
                            print("Received audio data:", audio)
                            logger.info(f"Received audio data: {audio}")
                            # 发送数据给ESP客户端
                            await ws.send(audio)
                            # 进行后续的处理
                        except Exception as e:
                            print("Error receiving TTS:", e)
                            logger.error(f"Error receiving TTS: {e}")
    except Exception as e:
        print("Error in websocket connection:", e)
        logger.error(f"Error in websocket connection: {e}")





@message_bp.websocket("/draw")
async def draw(request: Request, ws: Websocket) -> None:
    # 调用生成图像的函数，获取生成的数据或文件名
    generated_data = du_draw()  # 假设 du_draw() 返回生成的数据或文件名
    if generated_data is None:
        return  # 如果生成数据为 None，则直接退出函数

    if isinstance(generated_data, str):  # 假设返回的是文件名
        try:
            async with aiofiles.open(generated_data, mode='rb') as f:
                image_data = await f.read()
        except FileNotFoundError:
            # 处理文件找不到的情况
            return
        await ws.send(str(generated_data))

        # 发送图像数据给 WebSocket 客户端
        # await ws.send(image_data)
    else:
        # 处理其他生成数据的情况，例如直接发送生成的数据
        await ws.send(str(generated_data))
